"""
使用streamlit 上传文档 然后基于RAG 检索文件
pip install streamlit==1.39.0
#TOML（Tom's Obvious, Minimal Language）是一种配置文件格式，旨在比 JSON 更加简洁，易于人类编辑，并且比 YAML 更加安全。
pip install toml
pip install langchain_chroma

运行 切换到该文件：
 streamlit run langchain\llm_langchain_10_RAG_查询上传文档中的内容.py
"""
import os
import tempfile

import streamlit as st
from dotenv import load_dotenv
from langchain.agents import create_react_agent, AgentExecutor
from langchain.memory import ConversationBufferMemory
from langchain.tools.retriever import create_retriever_tool
from langchain_chroma import Chroma
from langchain_community.callbacks.streamlit.streamlit_callback_handler import StreamlitCallbackHandler
from langchain_community.chat_message_histories import StreamlitChatMessageHistory
# 用于加载文本文件
from langchain_community.document_loaders import TextLoader
from langchain_core.prompts import PromptTemplate
# HuggingFace 模型嵌入
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_openai import ChatOpenAI
from langchain_text_splitters import RecursiveCharacterTextSplitter

# 设置huggingface 本地模型下载地址
# os.environ['HF_HOME'] = 'D:/huggingface models/hub'
# 设置国内 hf代理地址
# os.environ['HF_ENDPOINT'] = 'https://modelscope.cn'

# 记录历史聊天记录
# 加载txt的库
# 向量数据库
# 提示次模板
# 文档分割
# 流程基于agent自动完成
# agent结果动态呈现在streamlit中 大模型思考后 回调streamlitcallback

# 自动加载 .env 文件中的环境变量
load_dotenv()

# -------------------------------
# Streamlit 页面配置
# -------------------------------
# 设置Streamlit 应用的页面标题和布局
st.set_page_config(page_title="文档问答", layout="wide")
# 设置应用标题
st.title("文档问答")

# -------------------------------
# 上传文件
# -------------------------------
# 允许上传多个 txt 文件
uploaded_files = st.sidebar.file_uploader(
    label="上传txt文件", type=["txt"], accept_multiple_files=True
)

# 如果没有上传文件 提示用户上传文件并停止运行
if not uploaded_files:
    st.info("请先上传txt文档。")
    st.stop()


# -------------------------------
# 创建文档检索器函数
# -------------------------------
# @st.cache_resource 会缓存函数返回值，避免重复计算
@st.cache_resource(ttl="1h")
def configure_retirever(uploaded_files):
    """
       功能：
           1. 将上传的文件写入临时目录
           2. 加载文件内容
           3. 使用 RecursiveCharacterTextSplitter 分割文档
           4. 使用 HuggingFaceEmbeddings 生成向量
           5. 构建 Chroma 向量数据库
           6. 返回检索器 retriever
       """
    # 存储所有文档内容
    docs = []
    # 临时目录
    temp_dir = tempfile.TemporaryDirectory(dir=r"D:\\temp")

    for file in uploaded_files:
        # 将上传文件写入临时目录
        temp_file_path = os.path.join(temp_dir.name, file.name)
        with open(temp_file_path, "wb") as f:
            f.write(file.getvalue())

        # 使用 TextLoader 加载文本
        loader = TextLoader(temp_file_path, encoding="utf-8")
        # 读取文件内容并添加到 docs 列表
        docs.extend(loader.load())

    # -------------------------------
    # 文档分割（防止太大影响检索效果）
    # -------------------------------
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    splits = text_splitter.split_documents(docs)

    if not splits:
        st.error("文档分割失败，未能生成有效的文本片段！")
        st.stop()

    # -------------------------------
    # 嵌入向量生成
    # -------------------------------
    embeddings = HuggingFaceEmbeddings(
        model_name="D:\models\sentence-transformersall-MiniLM-L6-v2",  # 本地模型路径
        # model_name="D:\models\BAAIbge-base-zh-v1.5",  # 本地模型路径
        model_kwargs={'device': 'cpu'}  # 可换成 'cuda' 使用 GPU
    )

    # -------------------------------
    # 构建向量数据库
    # -------------------------------
    vectordb = Chroma.from_documents(splits, embeddings)

    # 创建检索器
    retriever = vectordb.as_retriever()

    return retriever


# -------------------------------
# 配置检索器
# -------------------------------
retriever = configure_retirever(uploaded_files)

# -------------------------------
# 初始化聊天记录
# -------------------------------
if "messages" not in st.session_state or st.sidebar.button("清空聊天记录"):
    st.session_state["messages"] = [{"role": "assistant", "content": "您好，我是文档问答助手"}]

# 显示历史聊天记录
for msg in st.session_state.messages:
    st.chat_message(msg["role"]).write(msg["content"])

# -------------------------------
# 创建检索工具（用于 Agent）
# -------------------------------
tool = create_retriever_tool(
    retriever,
    "文档检索",
    "用于检索用户提出的问题，并基于检索到的文档内容进行回复。"
)

tools = [tool]

# -------------------------------
# 创建聊天历史对象和内存
# -------------------------------
msgs = StreamlitChatMessageHistory()

# 创建对话缓冲区内存
memory = ConversationBufferMemory(
    chat_memory=msgs, return_messages=True, memory_key="chat_history", output_key="output"
)

# -------------------------------
# 创建提示词模板
# -------------------------------
instructions = """
你是一个设计用于查询文档来回答问题的代理。
你可以使用文档检索工具，并基于检索内容来回答问题。
你可以不查询文档就知道答案，但是你任然应该查询文档来获取答案。
如果你从文档中找不到任何信息用于回答问题，则只需要返回“抱歉，这个问题我还不知道”作为答案。
"""

# 基础提示词模板
base_prompt_template = """
{instructions}

TOOLS:
------

You have access to the following tools:

{tools}

To use a tool, please use the following format:

•```
Thought: Do I need to use a tool? Yes
Action: the action to take, should be one of [{tool_names}]
Action Input: the input to the action
Observation: the result of the action
•```

When you have a response to say to the Human, or if you do not need to use a tool, you MUST use the format:

•```
Thought: Do I need to use a tool? No
Final Answer: [your response here]
•```

Begin!

Previous conversation history:
{chat_history}

New input: {input}
{agent_scratchpad}"""

# 创建基础提示词模板
base_prompt = PromptTemplate.from_template(base_prompt_template)

# 创建部分填充提示词模板
prompt = base_prompt.partial(instructions=instructions)

# -------------------------------
# 创建大模型
# -------------------------------
llm = ChatOpenAI(api_key=os.getenv("XUNFEI_API_KEY"), openai_api_base=os.getenv("XUNFEI_API_BASE"), model="xdeepseekv32exp", streaming=True, temperature=0.7)

# -------------------------------
# 创建 Agent 和执行器
# -------------------------------
agent = create_react_agent(llm, tools, prompt)

# 创建Agent执行器
agent_executor = AgentExecutor(agent=agent,
                               tools=tools,
                               memory=memory,
                               verbose=True,
                               handle_parsing_errors="没有从知识库中检索到相似内容")

# -------------------------------
# 用户输入处理
# -------------------------------
user_query = st.chat_input(placeholder="请开始提问吧")

# 如果有用户输入的查询
if user_query:
    # 保存用户消息
    st.session_state.messages.append({"role": "user", "content": user_query})
    # 显示用户消息
    st.chat_message("user").write(user_query)

    with st.chat_message("assistant"):
        # 创建 Streamlit 回调对象
        st_call_back = StreamlitCallbackHandler(st.container())
        # agent 执行过程日志回调显示 Streamlit Container （如思考，选择工具，执行拆线呢，观察结果等）
        config = {"callbacks": [st_call_back]}

        # 执行 Agent
        response = agent_executor.invoke({"input": user_query}, config=config)

        # 保存助手消息
        st.session_state.messages.append({"role": "assistant", "content": response["output"]})
        # 显示助手消息
        st.write(response["output"])
